package com.example.emailWeb.ntru.encrypt;

import com.example.emailWeb.ntru.polynomial.DenseTernaryPolynomial;
import com.example.emailWeb.ntru.polynomial.IntegerPolynomial;
import com.example.emailWeb.ntru.polynomial.LongPolynomial;

import java.nio.ByteBuffer;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Random;

public class NtruProxyUserA {
    int N,p,q,maxM1;
    boolean sparse;
    NtruProxyParameters params;
    public IntegerPolynomial fA,invertfA,gA;
    IntegerPolynomial PKA;
    public IntegerPolynomial r1,r2,e,a;

    private void init(){
        this.N = params.N;
        this.p = params.p;
        this.q = params.q;
        this.maxM1 = params.maxM1;
        this.sparse = params.sparse;
        this.gA = generateG(new SecureRandom());
        generateFa();
        generatePKA();
        r1 = generateR(); //发给C
        e = generateR();
        a = generateA(); // 发给B
    }

    public NtruProxyUserA(NtruProxyParameters param) {
        this.params = param;
        init();
    }

    public NtruProxyUserA() {
        this.params = NtruProxyParameters.APR2011_439;
        init();
    }

    public NtruProxyUserA(IntegerPolynomial PKA) {
        this.params = NtruProxyParameters.APR2011_439;
        init();
        this.PKA = PKA;
    }

    /**
     * fA * invertfA(mod q)= 1;
     * fA(mod p)= 1;
     */
    public void generateFa(){
        int[] coeffs = new int[N];
        int numP = N / 10;
//        int numP = 3;
        int Count = numP;
//        int Count = 10;
        Random random = new Random();
        while(true) {
//            while (Count > 0) {
//                int Index = random.nextInt(N);
//                if (coeffs[Index] == 0) {
//                    coeffs[Index] = p;
//                    Count -= 1;
//                }
//            }
            coeffs[0] = 1;
            fA = new IntegerPolynomial(coeffs);
            if (fA.isInvertiblePow2()) {
                IntegerPolynomial invertfA = fA.invertFq(q);
                if (invertfA != null) {
                    this.invertfA = invertfA;
                    break;
                }
            }
        }
    }

    private IntegerPolynomial generateG(Random rng) {
        final int N = params.N;
        int dg = params.dg; // dg = N / 3;

        while (true) {
            DenseTernaryPolynomial g = DenseTernaryPolynomial.generateRandom(N, dg, dg-1, rng);
            //产生N个数，dg个1，dg-1个-1，其余的为0，rng用作shuffle的参数。
            if (g.isInvertiblePow2())
                return g;
        }
    }

    /**
     * 生成公钥PKA;
     * PKA = p * g * invertfA(mod q);
     * @return
     */
    private void generatePKA(){
        IntegerPolynomial g = this.gA.clone();
        g.mult3(q); //p*gA
        PKA = g;
        PKA = PKA.mult(invertfA,q);
    }

    public IntegerPolynomial generateR() {
        Random rng = new SecureRandom();
        final int N = params.N;
        int dg = params.dg; // dg = N / 3;

        while (true) {
            DenseTernaryPolynomial g = DenseTernaryPolynomial.generateRandom(N, dg, dg-1, rng);
            //产生N个数，dg个1，dg-1个-1，其余的为0，rng用作shuffle的参数。
            if (g.isInvertiblePow2())
                return g;
        }
    }
    public IntegerPolynomial generateA(){
        IntegerPolynomial A = e.clone();//避免把e直接修改掉，如果不clone(),A指向e的引用，A和e会同步修改。
        A.mult3(q);//e*p
        A.add(fA,q);//e*p+fA
        A = A.mult(r1,q);
        return A;
    }

    public IntegerPolynomial getPKA() {
        return PKA;
    }

    public IntegerPolynomial generateCA(String mStr){
        byte[] M = mStr.getBytes();
        System.out.println("明文字节: "+ Arrays.toString(M));
        return generateCA(M);
    }

    public IntegerPolynomial generateCA(byte[] M){
        //R2是什么不重要，R2的值不需要存放，因为在解密的过程中R2被消掉了
        IntegerPolynomial R2 = generateR(); //R2
        this.r2 = R2.clone();
        IntegerPolynomial res1 = R2.mult(PKA); // R2*PKA

        ByteBuffer mBuf = ByteBuffer.allocate(M.length+3);
        mBuf.put((byte)M.length);
        mBuf.put(M);
        mBuf.put((byte)0);//在字节的尾部填充，填充什么都可以。
        mBuf.put((byte)0);//在字节数组的尾部填充，避免从字节数组恢复成字符串的时候，出现乱码
        byte[] cM = mBuf.array();

        IntegerPolynomial mTrin = IntegerPolynomial.fromBinary3Sves(cM, N, maxM1>0);   // don't use the constant coeff if maxM1 is set; see below
        System.out.println("明文的三元多项式：");
        System.out.println(mTrin);
        res1.add(mTrin);
        return res1;
    }

    public LongPolynomial generateCALong(byte[] M){
        IntegerPolynomial R2 = generateR(); //R2
        this.r2 = R2.clone();
        LongPolynomial R2Long = R2.clone().toLongPolynomial();
        System.out.println("r2 in generateCA: "+R2);
        LongPolynomial res1 = R2Long.mult(PKA.toLongPolynomial()); // R2*PKA
        IntegerPolynomial mTrin = IntegerPolynomial.fromBinary3Sves(M, N, maxM1>0);   // don't use the constant coeff if maxM1 is set; see below
        System.out.println("明文的三元多项式：");
        System.out.println(mTrin);
        LongPolynomial mTrinLong = mTrin.toLongPolynomial();
        res1.add(mTrinLong);
        return res1;
//        byte[] C = res1.toBinary(q);
//        return res1;
    }

    @Override
    public String toString() {
        return "NtruProxyUserA{" +
                "fA=" + fA +
                ", PKA=" + PKA +
                ", r1=" + r1 +
                ", r2=" + r2 +
                ", a=" + a +
                '}';
    }
}
